package cn.taskservice.common.utils.thread;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.concurrent.*;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;

/**
 * 多任务并发统计模板类，支持任务级超时控制、降级处理、日志记录、执行耗时监控等能力。
 */
@Slf4j
@Component
public class StatAggregationTemplate {

    @Autowired
    @Qualifier("ioTaskExecutor") // 指定线程池
    private ThreadPoolTaskExecutor statThreadPool;

    private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();

    /**
     * 聚合执行带名称的统计任务，每个任务支持超时、异常降级、耗时打印。
     *
     * @param taskMap          任务名称与任务定义的映射
     * @param perTaskTimeoutMs 单个任务的超时时间（毫秒）
     * @param overallTimeoutMs 整体聚合超时时间（毫秒）
     * @param <T>              任务结果类型
     * @return 统计结果 Map
     */
    public <T> Map<String, T> aggregate(Map<String, TaskWrapper<T>> taskMap, long perTaskTimeoutMs, long overallTimeoutMs) {
        Map<String, CompletableFuture<T>> futureMap = new LinkedHashMap<>();

        for (Map.Entry<String, TaskWrapper<T>> entry : taskMap.entrySet()) {
            String name = entry.getKey();
            TaskWrapper<T> wrapper = entry.getValue();
            futureMap.put(name, wrapWithTimeoutAndFallback(name, wrapper, perTaskTimeoutMs));
        }

        CompletableFuture<Void> allDone = CompletableFuture.allOf(futureMap.values().toArray(new CompletableFuture[0]));

        try {
            allDone.get(overallTimeoutMs, TimeUnit.MILLISECONDS);
        } catch (TimeoutException e) {
            log.warn("[大屏聚合统计] 整体执行超时 {}ms", overallTimeoutMs);
        } catch (Exception e) {
            log.error("[大屏聚合统计] 执行异常", e);
        }

        return futureMap.entrySet().stream().collect(Collectors.toMap(
                Map.Entry::getKey,
                e -> {
                    try {
                        return e.getValue().join(); // 使用 join 获取结果，抛出 CompletionException
                    } catch (CompletionException ex) {
                        log.error("[{}] 任务执行异常: {}", e.getKey(),
                                ex.getCause() != null ? ex.getCause().getMessage() : ex.getMessage(), ex);
                        return null; // 降级返回 null
                    }
                },
                (a, b) -> b,
                LinkedHashMap::new
        ));
    }

    /**
     * 包装单个任务，加入线程池异步执行、耗时统计、异常处理、超时控制、降级处理。(线程池为自定义的IO密集型配置)
     */
    private <T> CompletableFuture<T> wrapWithTimeoutAndFallback(String name, TaskWrapper<T> taskWrapper, long timeoutMs) {
        CompletableFuture<T> main = CompletableFuture.supplyAsync(() -> {
            long start = System.currentTimeMillis();
            try {
                T result = taskWrapper.getTask().get();
                long duration = System.currentTimeMillis() - start;
                log.info("[{}] 执行成功，耗时: {}ms", name, duration);
                return result;
            } catch (Exception e) {
                log.error("[{}] 执行异常: {}", name, e.getMessage(), e);
                throw e;
            }
        }, statThreadPool);

        CompletableFuture<T> timeoutFuture = failAfter(timeoutMs);

        return main.applyToEither(timeoutFuture, Function.identity())
                .exceptionally(ex -> {
                    log.warn("[{}] 异常或超时，将执行降级逻辑", name);
                    return taskWrapper.getFallback().apply(ex);
                });
    }

    /**
     * 返回一个在指定时间后失败的 CompletableFuture。
     */
    private <T> CompletableFuture<T> failAfter(long timeoutMs) {
        CompletableFuture<T> promise = new CompletableFuture<>();
        scheduler.schedule(() -> promise.completeExceptionally(new TimeoutException("任务超时")), timeoutMs, TimeUnit.MILLISECONDS);
        return promise;
    }

    /**
     * 封装任务和降级处理器
     */
    public static class TaskWrapper<T> {
        private final Supplier<T> task;
        private final Function<Throwable, T> fallback;

        public TaskWrapper(Supplier<T> task, Function<Throwable, T> fallback) {
            this.task = task;
            this.fallback = fallback;
        }

        public Supplier<T> getTask() {
            return task;
        }

        public Function<Throwable, T> getFallback() {
            return fallback;
        }

        public static <T> TaskWrapper<T> of(Supplier<T> task, Function<Throwable, T> fallback) {
            return new TaskWrapper<>(task, fallback);
        }
    }
}
